from collections import deque
import xxhash
import numpy as np

from nanovllm.engine.sequence import Sequence


class Block:
    """
    实现了一种高效的分页式内存管理机制
    """
    def __init__(self, block_id):
        # 块的唯一标识符
        self.block_id = block_id
        # 引用计数，用于追踪有多少个序列正在使用这个块
        self.ref_count = 0
        # 块内容的哈希值，用于实现前缀缓存
        self.hash = -1
        # 块中存储的 token ID 序列
        self.token_ids = []

    def update(self, hash: int, token_ids: list[int]):
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []


class BlockManager:
    """
    PagedAttention（分页注意力机制）
    分页注意力机制是一种用于处理长序列的注意力机制，它将输入序列分成多个块（block），每个块的长度为固定的大小（block_size）。每个块都有一个唯一的标识符（block_id），用于在注意力计算中引用。
    分页注意力机制的主要优势是可以处理长序列，而不需要计算所有 token 之间的注意力。它通过将序列分成多个块，只计算每个块内部的注意力，从而减少了计算复杂度。
    blocks 数组：存储所有 Block 实例的主数组
    hash_to_block_id 映射：实现缓存复用的哈希表
    free_block_ids 队列：管理可用的空闲块
    used_block_ids 集合：跟踪正在使用的块
    """
    def __init__(self, num_blocks: int, block_size: int):
        # 每个块的大小（能存储的 token 数量）
        self.block_size = block_size 
        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]
        # 哈希值到块 ID 的映射，用于快速查找相同内容的块
        self.hash_to_block_id: dict[int, int] = dict()
        # 空闲块 ID 的队列
        self.free_block_ids: deque[int] = deque(range(num_blocks))
        self.used_block_ids: set[int] = set()

    @classmethod
    #计算 token 序列的哈希值
    def compute_hash(cls, token_ids: list[int], prefix: int = -1):
        h = xxhash.xxh64()
        if prefix != -1:
            h.update(prefix.to_bytes(8, "little"))
        h.update(np.array(token_ids).tobytes())
        return h.intdigest()
    # 块分配
    def _allocate_block(self, block_id: int) -> Block:
        block = self.blocks[block_id]
        assert block.ref_count == 0
        block.reset()
        self.free_block_ids.remove(block_id)
        self.used_block_ids.add(block_id)
        return self.blocks[block_id]
    # 块释放
    def _deallocate_block(self, block_id: int) -> Block:
        assert self.blocks[block_id].ref_count == 0
        self.used_block_ids.remove(block_id)
        self.free_block_ids.append(block_id)
    # 检查是否有足够的空闲块来分配给序列    
    def can_allocate(self, seq: Sequence) -> bool:
        return len(self.free_block_ids) >= seq.num_blocks
    # 序列分配所需的块
    def allocate(self, seq: Sequence):
        # 确保序列尚未分配块
        assert not seq.block_table 
        h = -1
        # 缓存未命中标志
        cache_miss = False
        for i in range(seq.num_blocks):# 遍历序列需要的所有块
             # 获取当前块的token_ids
            token_ids = seq.block(i)
            # 计算当前块的哈希值，如果是完整块则基于前缀哈希计算，否则设为-1
            # 这里实现了前缀链式哈希计算
            # 每个完整块的哈希值基于前一个块的哈希值计算
            # 这样可以形成唯一标识整个前缀链的哈希序列
            h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1
            #根据哈希值查找相同内容的块
            block_id = self.hash_to_block_id.get(h, -1)
            # 检查是否缓存未命中（哈希不存在或内容不匹配）
            if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
                cache_miss = True
            if cache_miss: # 缓存未命中时，分配新块
                block_id = self.free_block_ids[0]
                block = self._allocate_block(block_id)
            else:
                # 缓存命中时，增加缓存命中的token数量
                seq.num_cached_tokens += self.block_size
                # 如果块正在被使用，增加引用计数；否则分配该块
                if block_id in self.used_block_ids:
                    # 增加引用计数，实现块共享
                    block = self.blocks[block_id]
                    block.ref_count += 1
                else:
                    block = self._allocate_block(block_id)
            if h != -1:# 如果哈希有效，更新块的哈希和内容，并更新哈希表
                block.update(h, token_ids)
                self.hash_to_block_id[h] = block_id
             # 将块ID添加到序列的块表中    
            seq.block_table.append(block_id)
    # 序列释放块
    def deallocate(self, seq: Sequence):
        for block_id in reversed(seq.block_table):
            block = self.blocks[block_id]
            block.ref_count -= 1
            if block.ref_count == 0:
                self._deallocate_block(block_id)
        seq.num_cached_tokens = 0
        seq.block_table.clear()

    def can_append(self, seq: Sequence) -> bool:
        return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

    def may_append(self, seq: Sequence):
        block_table = seq.block_table # 获取序列的块表
        last_block = self.blocks[block_table[-1]] # 获取最后一个块
        if len(seq) % self.block_size == 1: # 序列长度是块大小的整数倍+1，需要新块
            assert last_block.hash != -1
            # 分配新块
            block_id = self.free_block_ids[0] 
            self._allocate_block(block_id)
            # 添加新块到块表
            block_table.append(block_id)
        elif len(seq) % self.block_size == 0: # 序列长度是块大小的整数倍，当前块已填满
            assert last_block.hash == -1
             # 获取当前块的token_ids
            token_ids = seq.block(seq.num_blocks-1)
            # 获取前一个块的哈希作为前缀（如果存在）
            prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1
            # 基于前缀计算当前块的哈希
            h = self.compute_hash(token_ids, prefix)
            # 更新块的哈希和内容
            last_block.update(h, token_ids)
            # 更新哈希映射表，使该块可被其他序列复用
            self.hash_to_block_id[h] = last_block.block_id
        else: # 序列长度既不是块大小的整数倍，也不是整数倍+1
            assert last_block.hash == -1
